package com.peng.leetcode.heap;

import java.util.ArrayList;
import java.util.LinkedList;
import java.util.List;
import java.util.PriorityQueue;

/**
 * KSmallestPairs
 *
 * @author: lupeng6
 * @create: 2020/12/11 19:21
 */
public class KSmallestPairs {

    public static void main(String[] args) {

    }

    static class Solution {
        public List<List<Integer>> kSmallestPairs(int[] nums1, int[] nums2, int k) {
            PriorityQueue<List<Integer>> queue = new PriorityQueue<>(k, (a, b) -> (b.get(0) + b.get(1)) - (a.get(0) + a.get(1)));
            for (int i = 0; i < nums1.length; i++) {
                for (int j = 0; j < nums2.length; j++) {
                    if (queue.size() == k && (queue.peek().get(0) + queue.peek().get(1)) < nums1[i] + nums2[j]) {
                        break;
                    }

                    if (queue.size() == k) {
                        queue.poll();
                    }
                    List<Integer> node = new ArrayList<>(2);
                    node.add(nums1[i]);
                    node.add(nums2[j]);
                    queue.add(node);
                }
            }

            List<List<Integer>> ret = new LinkedList<>();
            for (int i = 0; i < k && !queue.isEmpty(); i++) {
                ret.add(0, queue.poll());
            }
            return ret;
        }
    }
}
